use crate::types::{Position, PoolInfo, SubscribeMessage, WebSocketMessage};
use crate::utils::{TokenValidator, parse_pool_from_account_data};
use anyhow::Result;
use futures_util::{SinkExt, StreamExt};
use log::{error, info, warn};
use solana_client::rpc_client::RpcClient;
use solana_sdk::{
    commitment_config::CommitmentConfig,
    pubkey::Pubkey,
    signature::{Keypair, Signature},
    signer::Signer,
    transaction::Transaction,
    instruction::Instruction,
    compute_budget::ComputeBudgetInstruction,
    commitment_config::CommitmentLevel,
};
use std::collections::HashMap;
use std::str::FromStr;
use std::time::{Duration, SystemTime};
use tokio::time::interval;
use tokio_tungstenite::{connect_async, tungstenite::Message};
use base64;
use solana_transaction_status::UiTransactionEncoding;
use solana_client::rpc_request::TokenAccountsFilter;
use base64::Engine;

/// CLMM狙击机器人结构体
/// 用于监控和交易 Raydium CLMM 池子
pub struct ClmmSniperBot {
    helius_ws_url: String,          // Helius WebSocket URL
    rpc_client: RpcClient,          // Solana RPC 客户端
    keypair: Keypair,               // 钱包密钥对
    positions: HashMap<String, Position>,  // 当前持仓信息
    token_validator: TokenValidator, // 代币验证器
    clmm_program_id: Pubkey,        // Raydium CLMM 程序 ID
    buy_amount_sol: f64,            // 每次买入的 SOL 数量
    min_liquidity_sol: f64,         // 最小流动性要求
    max_liquidity_sol: f64,         // 最大流动性要求
    sell_profit_percentage: f64,    // 目标利润百分比
    max_holding_time: u64,  // 最大持仓时间
    compute_unit_limit: u32,        // 计算单元限制
    compute_unit_price: u64,        // 计算单元价格
    stop_loss_percentage: f64,      // 止损百分比
}

impl ClmmSniperBot {
    /// 创建新的 CLMM 狙击机器人实例
    pub async fn new(
        helius_ws_url: String,
        helius_rpc_url: String,
        keypair: Keypair,
        buy_amount_sol: f64,
        min_liquidity_sol: f64,
        max_liquidity_sol: f64,
        sell_profit_percentage: f64,
        max_holding_time: u64,
        compute_unit_limit: u32,
        compute_unit_price: u64,
        stop_loss_percentage: f64,
    ) -> Result<Self> {
        // 初始化 RPC 客户端
        let rpc_client = RpcClient::new_with_commitment(
            helius_rpc_url.clone(),
            CommitmentConfig::confirmed(),
        );
        
        // 初始化代币验证器
        let token_validator = TokenValidator::new(helius_rpc_url);
        
        // Raydium CLMM 程序 ID
        let clmm_program_id = Pubkey::from_str("CAMMCzo5YL8w4VFF8KVHrK22GGUsp5VTaW7grrKgrWqK")?;
        
        Ok(Self {
            helius_ws_url,
            rpc_client,
            keypair,
            positions: HashMap::new(),
            token_validator,
            clmm_program_id,
            buy_amount_sol,
            min_liquidity_sol,
            max_liquidity_sol,
            sell_profit_percentage,
            max_holding_time,
            compute_unit_limit,
            compute_unit_price,
            stop_loss_percentage,
        })
    }

    /// 启动机器人
    /// 同时启动 WebSocket 监控和定时卖出检查
    pub async fn start(&mut self) -> Result<()> {
        // 使用 unsafe 块来避免多次可变借用 self
        let ws_fut = {
            let this = self as *mut Self;
            async move { unsafe { &mut *this }.start_websocket_monitoring().await }
        };
        let sell_fut = {
            let this = self as *mut Self;
            async move { unsafe { &mut *this }.start_sell_timer().await }
        };
        tokio::try_join!(ws_fut, sell_fut)?;
        Ok(())
    }

    /// 启动 WebSocket 监控
    /// 用于监听新池子的创建
    async fn start_websocket_monitoring(&mut self) -> Result<()> {
        // 建立 WebSocket 连接
        info!("正在连接到 Helius WebSocket: {}", self.helius_ws_url);
        let (ws_stream, _) = connect_async(&self.helius_ws_url).await?;
        info!("WebSocket 连接成功建立");
        let (mut ws_sender, mut ws_receiver) = ws_stream.split();

        // 订阅 CLMM 程序账户变化
        let subscribe_msg = SubscribeMessage {
            jsonrpc: "2.0".to_string(),
            id: 1,
            method: "programSubscribe".to_string(),
            params: vec![
                serde_json::Value::String(self.clmm_program_id.to_string()),
                serde_json::json!({
                    "commitment": "confirmed",
                    "encoding": "base64",
                    "filters": []
                })
            ],
        };

        // 发送订阅消息
        let subscribe_text = serde_json::to_string(&subscribe_msg)?;
        info!("正在发送订阅消息: {}", subscribe_text);
        ws_sender.send(Message::Text(subscribe_text)).await?;
        info!("订阅消息已发送，等待确认...");
        
        // 等待订阅确认
        while let Some(msg) = ws_receiver.next().await {
            match msg? {
                Message::Text(text) => {
                    info!("收到WebSocket消息: {}", text);
                    if text.contains("subscription") {
                        info!("CLMM程序订阅已确认");
                        break;
                    }
                    if let Err(e) = self.handle_websocket_message(&text).await {
                        error!("处理WebSocket消息失败: {}", e);
                    }
                }
                Message::Close(_) => {
                    warn!("WebSocket连接已关闭");
                    break;
                }
                _ => {}
            }
        }

        Ok(())
    }

    /// 处理 WebSocket 消息
    async fn handle_websocket_message(&mut self, message: &str) -> Result<()> {
        info!("正在处理WebSocket消息: {}", message);
        let ws_msg: WebSocketMessage = serde_json::from_str(message)?;
        
        // 检查是否是程序通知消息
        if let Some(method) = &ws_msg.method {
            info!("收到方法调用: {}", method);
            if method == "programNotification" {
                info!("收到程序通知消息");
                if let Some(params) = &ws_msg.params {
                    info!("程序通知参数: {}", serde_json::to_string(params)?);
                    self.handle_program_notification(params).await?;
                }
            }
        }
        
        Ok(())
    }

    /// 处理程序通知
    /// 检查是否是新池子创建，如果是则执行买入操作
    async fn handle_program_notification(&mut self, params: &serde_json::Value) -> Result<()> {
        // 解析程序通知参数
        if let Some(result) = params.get("result") {
            if let Some(value) = result.get("value") {
                if let Some(account_info) = value.get("account") {
                    if let Some(pubkey_str) = result.get("pubkey") {
                        let pubkey = Pubkey::from_str(pubkey_str.as_str().unwrap())?;
                        
                        info!("收到账户更新通知: {}", pubkey);
                        
                        // 检查是否是新创建的池子
                        match self.is_new_pool_creation(&account_info).await {
                            Ok(true) => {
                                info!("检测到新池子创建: {}", pubkey);
                                self.handle_new_pool_creation(pubkey).await?;
                            }
                            Ok(false) => {
                                info!("账户 {} 不是新池子创建", pubkey);
                            }
                            Err(e) => {
                                error!("检查池子创建失败: {}", e);
                            }
                        }
                    }
                }
            }
        }
        
        Ok(())
    }

    /// 检查是否是新池子创建
    /// 通过检查账户数据和所有者来验证
    async fn is_new_pool_creation(&self, account_info: &serde_json::Value) -> Result<bool> {
        // 检查账户数据是否存在
        if let Some(data) = account_info.get("data") {
            if let Some(data_str) = data.as_str() {
                info!("检查账户数据: {}", data_str);
                
                // 解码 base64 数据
                let data_bytes = base64::engine::general_purpose::STANDARD.decode(data_str)?;
                
                // 检查数据长度
                if data_bytes.len() < 8 {
                    info!("数据长度不足8字节");
                    return Ok(false);
                }
                
                // 检查账户标识符
                let account_discriminator = &data_bytes[0..8];
                let expected_discriminator = [0, 0, 0, 0, 0, 0, 0, 0]; // Raydium CLMM 池子的标识符
                
                info!("账户标识符: {:?}", account_discriminator);
                
                // 如果是新创建的池子，标识符应该匹配
                if account_discriminator == expected_discriminator {
                    // 检查账户所有者是否为 Raydium CLMM 程序
                    if let Some(owner) = account_info.get("owner") {
                        if let Some(owner_str) = owner.as_str() {
                            let owner_pubkey = Pubkey::from_str(owner_str)?;
                            info!("账户所有者: {}", owner_pubkey);
                            if owner_pubkey == self.clmm_program_id {
                                info!("确认是 CLMM 程序创建的账户");
                                return Ok(true);
                            }
                        }
                    }
                }
            }
        }
        
        Ok(false)
    }

    /// 处理新池子创建
    /// 验证代币条件并执行买入操作
    async fn handle_new_pool_creation(&mut self, pool_pubkey: Pubkey) -> Result<()> {
        // 获取池子账户数据
        let account_info = self.rpc_client.get_account(&pool_pubkey)?;
        let pool_info = parse_pool_from_account_data(&account_info.data, &pool_pubkey)?;
        
        // 验证代币条件
        let validation = self.token_validator.validate_token(&pool_info).await?;
        
        // 检查是否满足所有条件
        if validation.is_mint_renounced 
            && validation.is_not_freezable 
            && validation.liquidity_in_range {
            
            info!("池子 {} 满足所有条件，执行买入", pool_pubkey);
            self.execute_buy(&pool_info).await?;
        } else {
            info!("池子 {} 不满足条件，跳过", pool_pubkey);
        }
        
        Ok(())
    }

    /// 执行买入操作
    /// 构建并发送买入交易
    async fn execute_buy(&mut self, pool_info: &PoolInfo) -> Result<()> {
        // 构建买入交易
        let transaction = self.build_buy_transaction(pool_info).await?;
        
        // 发送交易
        let signature = self.rpc_client.send_transaction_with_config(
            &transaction,
            solana_client::rpc_config::RpcSendTransactionConfig {
                skip_preflight: true,
                preflight_commitment: Some(CommitmentLevel::Confirmed),
                encoding: Some(UiTransactionEncoding::Base64),
                max_retries: Some(3),
                min_context_slot: None,
            },
        )?;
        
        info!("买入交易已发送: {}", signature);
        
        // 等待交易确认
        let blockhash = self.rpc_client.get_latest_blockhash()?;
        self.rpc_client.confirm_transaction_with_spinner(
            &signature,
            &blockhash,
            CommitmentConfig::confirmed(),
        )?;
        
        // 记录持仓信息
        self.record_position(pool_info, signature).await?;
        Ok(())
    }

    /// 构建买入交易
    /// 包含计算单元限制和价格设置
    async fn build_buy_transaction(&self, pool_info: &PoolInfo) -> Result<Transaction> {
        let recent_blockhash = self.rpc_client.get_latest_blockhash()?;
        
        // 构建 Raydium CLMM 交换指令
        let swap_instruction = self.build_swap_instruction(pool_info)?;
        
        // 添加计算单元限制和价格指令
        let compute_budget_instruction = ComputeBudgetInstruction::set_compute_unit_limit(self.compute_unit_limit);
        let compute_price_instruction = ComputeBudgetInstruction::set_compute_unit_price(self.compute_unit_price);
        
        // 创建并签名交易
        let transaction = Transaction::new_signed_with_payer(
            &[compute_budget_instruction, compute_price_instruction, swap_instruction],
            Some(&self.keypair.pubkey()),
            &[&self.keypair],
            recent_blockhash,
        );
        
        Ok(transaction)
    }

    fn build_swap_instruction(&self, pool_info: &PoolInfo) -> Result<Instruction> {
        let user_token_account = self.get_token_account(&pool_info.mint_a)?;
        let user_wsol_account = self.get_token_account(&pool_info.mint_b)?;
        let accounts = vec![
            solana_sdk::instruction::AccountMeta::new(self.clmm_program_id, false),
            solana_sdk::instruction::AccountMeta::new(pool_info.id, false),
            solana_sdk::instruction::AccountMeta::new(user_token_account, false),
            solana_sdk::instruction::AccountMeta::new(user_wsol_account, false),
            solana_sdk::instruction::AccountMeta::new(self.keypair.pubkey(), true),
        ];
        let amount_in = (self.buy_amount_sol * 1e9) as u64;
        let mut data = vec![0u8];
        data.extend_from_slice(&amount_in.to_le_bytes());
        Ok(Instruction {
            program_id: self.clmm_program_id,
            accounts,
            data,
        })
    }

    fn get_token_account(&self, mint: &Pubkey) -> Result<Pubkey> {
        let token_accounts = self.rpc_client.get_token_accounts_by_owner(
            &self.keypair.pubkey(),
            TokenAccountsFilter::Mint(*mint),
        )?;
        if token_accounts.is_empty() {
            return Err(anyhow::anyhow!("未找到代币账户"));
        }
        Ok(Pubkey::from_str(&token_accounts[0].pubkey)?)
    }

    /// 记录持仓信息
    /// 在成功买入后记录交易详情
    async fn record_position(&mut self, pool_info: &PoolInfo, signature: Signature) -> Result<()> {
        // 等待交易确认
        let blockhash = self.rpc_client.get_latest_blockhash()?;
        self.rpc_client.confirm_transaction_with_spinner(
            &signature,
            &blockhash,
            CommitmentConfig::confirmed(),
        )?;
        
        // 获取代币账户余额
        let token_account = self.get_token_account(&pool_info.mint_a)?;
        let token_balance = self.rpc_client.get_token_account_balance(&token_account)?
            .ui_amount
            .unwrap_or(0.0) as u64;
            
        // 获取当前价格
        let current_price = self.get_current_price(&Position {
            pool_id: pool_info.id,
            token_mint: pool_info.mint_a,
            buy_time: SystemTime::now(),
            buy_amount: 0,
            token_balance: 0,
            buy_price: 0.0,
        }).await?;
        
        // 创建持仓记录
        let position = Position {
            pool_id: pool_info.id,
            token_mint: pool_info.mint_a,
            buy_time: SystemTime::now(),
            buy_amount: (self.buy_amount_sol * 1e9) as u64,
            token_balance,
            buy_price: current_price,
        };
        
        info!(
            "记录新持仓: 池子={}, 代币={}, 买入金额={} SOL, 代币余额={}, 买入价格={}",
            pool_info.id,
            pool_info.mint_a,
            self.buy_amount_sol,
            token_balance,
            current_price
        );
        
        // 保存持仓信息
        self.positions.insert(pool_info.id.to_string(), position);
        Ok(())
    }

    /// 启动定时卖出检查
    /// 定期检查持仓是否需要卖出
    async fn start_sell_timer(&mut self) -> Result<()> {
        let mut interval = interval(Duration::from_secs(60));
        
        loop {
            interval.tick().await;
            self.check_and_sell_positions().await?;
        }
    }

    /// 检查并卖出持仓
    /// 遍历所有持仓，检查是否需要卖出
    async fn check_and_sell_positions(&mut self) -> Result<()> {
        for (pool_id, position) in self.positions.clone() {
            if self.should_sell_position(&position).await? {
                self.execute_sell(&position).await?;
                self.positions.remove(&pool_id);
            }
        }
        
        Ok(())
    }

    /// 检查是否应该卖出持仓
    /// 基于持仓时间、利润和止损条件
    async fn should_sell_position(&self, position: &Position) -> Result<bool> {
        // 检查持仓时间
        let holding_time = SystemTime::now()
            .duration_since(position.buy_time)?
            .as_secs() ;
        
        // 如果超过最大持仓时间，卖出
        if holding_time >= self.max_holding_time {
            info!("持仓时间超过 {} 秒，准备卖出", self.max_holding_time);
            return Ok(true);
        }
        
        // 获取当前价格
        let current_price = self.get_current_price(position).await?;
        
        // 计算价格变化百分比
        let price_change_percentage = ((current_price - position.buy_price) / position.buy_price) * 100.0;
        
        // 检查止损条件
        if price_change_percentage <= -self.stop_loss_percentage {
            info!("触发止损，价格下跌 {}%", price_change_percentage.abs());
            return Ok(true);
        }
        
        // 检查止盈条件
        if price_change_percentage >= self.sell_profit_percentage {
            info!("达到目标利润 {}%，准备卖出", price_change_percentage);
            return Ok(true);
        }
        
        Ok(false)
    }

    /// 执行卖出操作
    /// 构建并发送卖出交易
    async fn execute_sell(&self, position: &Position) -> Result<()> {
        // 获取池子信息
        let pool_account = self.rpc_client.get_account(&position.pool_id)?;
        let pool_info = parse_pool_from_account_data(&pool_account.data, &position.pool_id)?;
        
        // 获取最新区块哈希
        let recent_blockhash = self.rpc_client.get_latest_blockhash()?;
        
        // 构建卖出指令
        let sell_instruction = self.build_sell_instruction(&pool_info, position)?;
        
        // 添加计算单元限制和价格指令
        let compute_budget_instruction = ComputeBudgetInstruction::set_compute_unit_limit(self.compute_unit_limit);
        let compute_price_instruction = ComputeBudgetInstruction::set_compute_unit_price(self.compute_unit_price);
        
        // 创建并签名交易
        let transaction = Transaction::new_signed_with_payer(
            &[compute_budget_instruction, compute_price_instruction, sell_instruction],
            Some(&self.keypair.pubkey()),
            &[&self.keypair],
            recent_blockhash,
        );
        
        // 发送交易
        let signature = self.rpc_client.send_transaction_with_config(
            &transaction,
            solana_client::rpc_config::RpcSendTransactionConfig {
                skip_preflight: true,
                preflight_commitment: Some(CommitmentLevel::Confirmed),
                encoding: Some(UiTransactionEncoding::Base64),
                max_retries: Some(3),
                min_context_slot: None,
            },
        )?;
        
        info!("卖出交易已发送: {}", signature);
        
        // 等待交易确认
        let blockhash = self.rpc_client.get_latest_blockhash()?;
        self.rpc_client.confirm_transaction_with_spinner(
            &signature,
            &blockhash,
            CommitmentConfig::confirmed(),
        )?;
        
        info!("卖出交易已确认: {}", signature);
        Ok(())
    }

    fn build_sell_instruction(&self, pool_info: &PoolInfo, position: &Position) -> Result<Instruction> {
        let user_token_account = self.get_token_account(&position.token_mint)?;
        let user_wsol_account = self.get_token_account(&pool_info.mint_b)?;
        let accounts = vec![
            solana_sdk::instruction::AccountMeta::new(self.clmm_program_id, false),
            solana_sdk::instruction::AccountMeta::new(pool_info.id, false),
            solana_sdk::instruction::AccountMeta::new(user_token_account, false),
            solana_sdk::instruction::AccountMeta::new(user_wsol_account, false),
            solana_sdk::instruction::AccountMeta::new(self.keypair.pubkey(), true),
        ];
        let amount_in = position.token_balance;
        let mut data = vec![1u8];
        data.extend_from_slice(&amount_in.to_le_bytes());
        Ok(Instruction {
            program_id: self.clmm_program_id,
            accounts,
            data,
        })
    }

    /// 获取当前价格
    /// 基于池子中的代币余额计算
    async fn get_current_price(&self, position: &Position) -> Result<f64> {
        // 获取池子信息
        let pool_account = self.rpc_client.get_account(&position.pool_id)?;
        let pool_info = parse_pool_from_account_data(&pool_account.data, &position.pool_id)?;
        
        // 获取代币账户余额
        let token_account = self.rpc_client.get_token_account_balance(&position.token_mint)?
            .ui_amount
            .unwrap_or(0.0);
        
        // 获取 WSOL 账户余额
        let wsol_account = self.get_token_account(&pool_info.mint_b)?;
        let wsol_balance = self.rpc_client.get_token_account_balance(&wsol_account)?
            .ui_amount
            .unwrap_or(0.0);
        
        // 计算当前价格
        // 使用 Raydium CLMM 的价格计算公式
        // 价格 = (WSOL余额 / 代币余额) * (1 + 手续费)
        let fee_rate = 0.0025; // 0.25% 手续费
        let price = (wsol_balance / token_account) * (1.0 + fee_rate);
        
        // 验证价格有效性
        if price <= 0.0 || price.is_infinite() || price.is_nan() {
            return Err(anyhow::anyhow!("无效的价格计算结果"));
        }
        
        Ok(price)
    }
} 